import sensor, image, lcd, time
import KPU as kpu
from machine import UART
import gc, sys
from fpioa_manager import fm


input_size = (224, 224)
labels = ['y']
anchors = [0.72, 0.59, 1.13, 1.16, 0.5, 0.34, 1.28, 0.63, 2.38, 2.19]

def lcd_show_except(e):
    import uio
    err_str = uio.StringIO()
    sys.print_exception(e, err_str)
    err_str = err_str.getvalue()
    img = image.Image(size=input_size)
    img.draw_string(0, 10, err_str, scale=1, color=(0xff,0x00,0x00))
    lcd.display(img)


class Comm:
    def __init__(self, uart):
        self.uart = uart

    def send_detect_result(self, objects, labels):

        self.uart.write("Y")



def init_uart():
    fm.register(7, fm.fpioa.UART1_TX, force=True)
    fm.register(6, fm.fpioa.UART1_RX, force=True)
    uart = UART(UART.UART1, 9600, 8, 0, 1, timeout=1000, read_buf_len=1)
    return uart




def main(anchors, labels, model_addr="model-153566.kmodel", sensor_window=input_size, lcd_rotation=0):
    uart = init_uart()
    datayes='N'
    datano='Y'
    flag=1

    while(True):
        data1=uart.readline()
        if data1 is not None:
            datayes= data1.decode('utf-8')
        time.sleep(1)

        if datayes=='Y':
            datayes='N'
            sensor.reset()

            sensor.set_pixformat(sensor.RGB565)
            sensor.set_framesize(sensor.QVGA)
            sensor.set_windowing(sensor_window)
            sensor.run(1)
            lcd.rotation(lcd_rotation)
            comm = Comm(uart)
            num=50
            num2=20
            sleeptime=8
            yes=1
            try:
                task = kpu.load(model_addr)
                kpu.init_yolo2(task, 0.5, 0.3, 5, anchors)
                while(yes):
                    while(num>0):
                        img = sensor.snapshot()
                        t = time.ticks_ms()
                        objects = kpu.run_yolo2(task, img)
                        t = time.ticks_ms() - t

                        if objects:
                            for obj in objects:
                                pos = obj.rect()
                                img.draw_rectangle(pos)
                                img.draw_string(pos[0], pos[1], "%s : %.2f" % (labels[obj.classid()], obj.value()), scale=2, color=(255, 0, 0))
                            comm.send_detect_result(objects, labels)
                            num=1
                            flag=0
                        img.draw_string(0, 200, "t:%dms" % (t), scale=2, color=(255, 0, 0))
                        img.draw_string(0, 2, "Upgrade to MaixCAM to use YOLOv8", scale=1.2, color=(255, 0, 0))
                        img.draw_string(0, 30, "wiki.sipeed.com/maixcam", scale=1.2, color=(255, 0, 0))
                        lcd.display(img)
                        num=num-1
                        if(flag==1 and num==0):
                            uart.write("N")


                        data2=uart.readline()
                        if data2 is not None:
                            datano=data2.decode('utf-8').strip()
                        if datano=='N':
                            yes=0
                            num=0
                            datano='Y'

                    time.sleep(sleeptime)
                    num=50
                    flag=1

            except Exception as e:
                sys.print_exception(e)
                lcd_show_except(e)

            finally:
                if task:
                    kpu.deinit(task)
                gc.collect()

        else:pass

if __name__ == "__main__":
    try:
        main(anchors=anchors, labels=labels, model_addr="model-153566.kmodel")
    except Exception as e:
        sys.print_exception(e)
        lcd_show_except(e)
    finally:
        gc.collect()


